
package cn.myeasyai.milvus;

import cn.myeasyai.config.MilvusConfig;
import cn.myeasyai.dto.UserDTO;
import cn.myeasyai.dto.UserInfo;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.InsertParam.Field;

import java.util.*;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;

/**
 * @author fushoujiang
 */
@Service
public class VectorService {
    @Resource
    private MilvusServiceClient milvusClient;

    @Resource
    private MilvusClientV2 milvusV2Client;

    @Resource
    private MilvusConfig milvusConfig;

    private static List<String> fieldList = Arrays.asList("face_url","face_feature_url","user_id","user_name","Auto_id");

    public List<List<SearchResp.SearchResult>> search(List<Float> floatList,Integer topK) {
        FloatVec floatVec = new FloatVec(floatList);

        SearchReq searchReq = SearchReq.builder()
                .collectionName(milvusConfig.getCollectionName())
                .data(Collections.singletonList(floatVec))
                .outputFields(fieldList)
                .topK(topK)
                .build();

        SearchResp searchResp = milvusV2Client.search(searchReq);
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        return searchResults;
    }
/*    public List<SearchResultsWrapper.IDScore> search(List<Float> floatList,Integer topK) {
        List<List<Float>> vectors = Collections.singletonList(floatList);
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(milvusConfig.getCollectionName())
                .withMetricType(MetricType.L2)
                .withTopK(topK)
                .withOutFields(fieldList)
                .withFloatVectors(vectors)
                .withVectorFieldName(milvusConfig.getVectorFieldName()).build();
        R<SearchResults> searchResults = this.milvusClient.search(searchParam);
        if (searchResults.getStatus() == 0) {
            return  new SearchResultsWrapper(searchResults.getData().getResults()).getIDScore(0);
        }
        return Collections.emptyList();
    }*/


    /**
     * 初始化商家向量数据
     * @param userDTO 用户信息
     * @param features 向量数据
     */
    public void initUserVector(UserDTO userDTO, List<Float> features) {
        Gson gson = new Gson();
        UserInfo userInfo = new UserInfo();
        userInfo.setVector(features);
        userInfo.setFace_url(userDTO.getFaceUrl());
        userInfo.setFace_feature_url(userDTO.getFaceFeatureUrl());
        userInfo.setUser_id(userDTO.getUserId());
        userInfo.setUser_name(userDTO.getUserName());

        String json = gson.toJson(userInfo);

        List<JsonObject> data = Arrays.asList(
                gson.fromJson(json, JsonObject.class)
        );

        InsertReq insertReq = InsertReq.builder()
                .collectionName(milvusConfig.getCollectionName())
                .data(data)
                .build();

        InsertResp insertResp = milvusV2Client.insert(insertReq);
        System.out.println(insertResp);

    }


    /**
     * 清空历史数据
     */
    public void cleanAll() {
        DeleteParam build = DeleteParam.newBuilder()
                .withCollectionName(milvusConfig.getCollectionName())
                .withExpr("Auto_id not in [1]")
                .build();
        this.milvusClient.delete(build);
    }

}
